#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import argparse
import torch
import sys
from transformer_denoising_model import PointCloudDenoisingDataset, PointCloudTransformerDenoiser, PointCloudDenoisingTrainer, PointCloudVisualizer
from torch.utils.data import DataLoader
from pathlib import Path

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="点云Transformer去噪模型训练脚本")
    
    # 数据集参数
    parser.add_argument('--noise_type', type=str, default='gaussian', help='噪声类型: gaussian, outlier, mixed')
    parser.add_argument('--noise_level', type=str, default='0.05', help='噪声级别')
    parser.add_argument('--data_dir', type=str, default='./data/modelnet_noisy', help='数据目录')
    parser.add_argument('--absolute_path', type=str, default=None, help='数据集的绝对路径，如果提供则优先使用')
    
    # 模型参数
    parser.add_argument('--embed_dim', type=int, default=128, help='嵌入维度')
    parser.add_argument('--depth', type=int, default=4, help='Transformer层数')
    parser.add_argument('--heads', type=int, default=4, help='注意力头数')
    
    # 训练参数
    parser.add_argument('--batch_size', type=int, default=16, help='批次大小')
    parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
    parser.add_argument('--lr', type=float, default=1e-4, help='学习率')
    
    # 其他参数
    parser.add_argument('--vis_samples', type=int, default=5, help='可视化样本数量')
    parser.add_argument('--cpu', action='store_true', help='强制使用CPU')
    parser.add_argument('--debug', action='store_true', help='启用调试模式，显示详细路径信息')
    
    args = parser.parse_args()
    
    # 打印当前工作目录和系统信息，帮助调试
    current_dir = os.getcwd()
    print(f"当前工作目录: {current_dir}")
    print(f"Python版本: {sys.version}")
    print(f"操作系统: {sys.platform}")
    
    # 构建数据集路径
    if args.absolute_path:
        # 使用提供的绝对路径
        dataset_dir = Path(args.absolute_path)
        print(f"使用提供的绝对路径: {dataset_dir}")
    else:
        # 使用相对路径
        data_dir = Path(args.data_dir)
        if not data_dir.is_absolute():
            data_dir = Path(current_dir) / data_dir
        dataset_dir = data_dir / f"{args.noise_type}_{args.noise_level}"
        print(f"使用构建的路径: {dataset_dir}")
    
    train_data_path = str(dataset_dir / "train_denoising.h5")
    test_data_path = str(dataset_dir / "test_denoising.h5")
    
    # 如果开启了调试模式，显示更多路径信息
    if args.debug:
        print("\n调试信息:")
        print(f"dataset_dir是否存在: {dataset_dir.exists()}")
        if dataset_dir.exists():
            print(f"dataset_dir内容: {[str(p) for p in dataset_dir.iterdir()]}")
        print(f"train_data_path是否存在: {os.path.exists(train_data_path)}")
        print(f"test_data_path是否存在: {os.path.exists(test_data_path)}")
    
    # 确保数据集存在
    if not os.path.exists(train_data_path) or not os.path.exists(test_data_path):
        print(f"错误: 找不到数据集 {train_data_path} 或 {test_data_path}")
        print("请尝试以下解决方案:")
        print("1. 运行 prepare_denoising_dataset.py 准备去噪数据集")
        print("2. 使用 --absolute_path 参数直接指定数据集所在的完整路径")
        print("3. 使用 --debug 参数启用调试模式，查看详细路径信息")
        print("4. 检查文件名大小写是否正确")
        exit(1)
    
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() and not args.cpu else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载数据集
    print(f"加载 {args.noise_type} 噪声 (级别: {args.noise_level}) 的数据集...")
    train_dataset = PointCloudDenoisingDataset(train_data_path)
    test_dataset = PointCloudDenoisingDataset(test_data_path)
    
    train_loader = DataLoader(
        train_dataset, 
        batch_size=args.batch_size, 
        shuffle=True, 
        num_workers=4,
        pin_memory=True if device.type == 'cuda' else False
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=args.batch_size, 
        shuffle=False, 
        num_workers=4,
        pin_memory=True if device.type == 'cuda' else False
    )
    
    # 创建模型
    model = PointCloudTransformerDenoiser(
        point_dim=3,
        embed_dim=args.embed_dim,
        depth=args.depth,
        heads=args.heads,
        mlp_dim=args.embed_dim * 2,
        dropout=0.1
    ).to(device)
    
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
    
    # 创建训练器
    trainer = PointCloudDenoisingTrainer(
        model=model,
        device=device,
        train_loader=train_loader,
        val_loader=test_loader,
        learning_rate=args.lr,
        weight_decay=1e-5
    )
    
    # 训练模型
    print(f"开始训练 {args.noise_type}_{args.noise_level} 去噪模型...")
    trainer.train(args.epochs)
    
    # 可视化结果
    print("\n开始可视化去噪结果...")
    visualizer = PointCloudVisualizer(model, device)
    vis_dir = Path('./results/visualization') / f"{args.noise_type}_{args.noise_level}"
    os.makedirs(str(vis_dir), exist_ok=True)  # 确保可视化目录存在
    visualizer.visualize_results(
        test_data_path,
        num_samples=args.vis_samples,
        save_dir=str(vis_dir)
    )
    
    print(f"训练和可视化完成! 结果保存在 {trainer.save_dir} 和 {vis_dir}") 